from functools import wraps
from flask import g, redirect, url_for


def login_required(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if hasattr(g, 'user'):
            return func(*args, **kwargs)
        else:
            return redirect(url_for("user.login"))

    return wrapper


def root_required(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if hasattr(g, 'user') and g.user.power == 0:
            return func(*args, **kwargs)
        else:
            return "权限不符"

    return wrapper


def teacher_required(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if hasattr(g, 'user') and g.user.power == 1:
            return func(*args, **kwargs)
        else:
            return "权限不符"

    return wrapper


def student_required(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if hasattr(g, 'user') and g.user.power == 2:
            return func(*args, **kwargs)
        else:
            return "权限不符"

    return wrapper
